import os, sys
import numpy as np
import re, math
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

CAP = 60000


class pts_fts_loader:
    def __init__(self, pts_file_, inte_file_, label_file_):
        self.pts_file = pts_file_
        self.inte_file = inte_file_
        self.label_file = label_file_
        self.pts_fts = []
        self.data_num = []
        self.label = []
        self.load_label()
        self.load_data()

    @staticmethod
    def readFile(filepath):
        with open(filepath, "r") as f:
            while (True):
                yield f.readline().strip()

    @staticmethod
    def filter_batch(pts_fts_, label_):
        label_ = np.expand_dims(label_, 1)
        pts_fts_label = np.concatenate((pts_fts_, label_), 1)
        rm_index = []
        idx = 0
        for pfl in pts_fts_label:
            if math.sqrt(pfl[0] * pfl[0] + pfl[1] * pfl[1] + pfl[2] * pfl[2]) >= 35:
                rm_index.append(idx)
            idx += 1
        return np.delete(pts_fts_label, rm_index, 0)

    def load_data(self):
        fts_of_file = []
        idx = 0
        for z in pts_fts_loader.readFile(os.path.join(self.inte_file)):
            if idx > CAP:
                break
            if z.strip() == "":
                break
            fts_of_file.append(float(z))
            idx += 1
        while idx < CAP:
            fts_of_file.append(float(0))
            idx += 1

        idx = 0
        file_pts_num = 0
        pts_of_file = []
        for z in pts_fts_loader.readFile(os.path.join(self.pts_file)):
            if idx > CAP:
                break
            if z.strip() == "":
                break
            try:
                node_pts = []
                li = [float(j) for j in re.split(',', z)]
                node_pts.append(li[0])
                node_pts.append(li[1])
                node_pts.append(li[2])
                file_pts_num += 1
                if pts_of_file == []:
                    pts_of_file = node_pts
                else:
                    pts_of_file = np.vstack((pts_of_file, node_pts))
            except Exception as ess:
                break
            idx += 1
        self.data_num.append(file_pts_num)
        while idx < CAP:
            pts_of_file = np.vstack((pts_of_file, [0, 0, 0]))
            idx += 1
        fts_of_file = np.expand_dims(fts_of_file, 1)
        self.pts_fts = np.concatenate((pts_of_file, fts_of_file), 1)

    def load_label(self):
        if self.label_file == None:
            return
        label_of_file = []
        idx = 0
        for z in pts_fts_loader.readFile(os.path.join(self.label_file)):
            if idx > CAP:
                break
            if z.strip() == "":
                break
            label_of_file.append(int(z))
            idx += 1
        while idx < CAP:
            label_of_file.append(int(0))
            idx += 1
        if self.label == []:
            self.label = label_of_file
        else:
            self.label = np.vstack((label_of_file, self.label))

    def get_pts_fts_label(self):
        return (self.pts_fts, np.array(self.data_num), self.label)


def visualization(data_train, label_train, i=0):
    # (batch, pointnum, pts_fts) (batch,label)
    def filter_batch(pts_fts_, label_):
        label_ = np.expand_dims(label_, 1)
        pts_fts_label = np.concatenate((pts_fts_, label_), 1)
        # pts_fts_label = np.squeeze(pts_fts_label,0)
        rm_index = []
        idx = 0
        all_pts_num = 0
        exclude_num = 0
        all_target_num = 0
        exclude_loaded_target_num = 0
        for pfl in pts_fts_label:
            if pfl[0] == pfl[1] == pfl[2] == pfl[3]:
                exclude_num += 1
                rm_index.append(idx)
                idx += 1
                continue
            all_pts_num += 1
            if pfl[4] != 0:
                all_target_num += 1
            if math.sqrt(pfl[0] * pfl[0] + pfl[1] * pfl[1] + pfl[2] * pfl[2]) >= 35 :
                rm_index.append(idx)
                exclude_num += 1
                if pfl[4] != 0:
                    exclude_loaded_target_num += 1
            idx += 1
        print("all pts:{} ,[ loaded pts:{} ],all target:{}, [loaded target:{} ],loaded target rate:{}".format(
            all_pts_num,
            len(pts_fts_label) - exclude_num,
            all_target_num,
            all_target_num - exclude_loaded_target_num,
            float(all_target_num - exclude_loaded_target_num) / float(all_target_num),
        ))
        return np.delete(pts_fts_label, rm_index, 0)

    def getcolor(var):
        return {
            0: 'c',  # DontCare      青色
            # 0: 'k',
            1: 'b',  # cyclist       蓝色
            2: 'g',  # tricycle      绿色
            3: 'r',  # smallMot      红色
            4: 'k',  # bigMot        黑色
            5: 'm',  # pedestrian    洋红
            6: 'y',  # crowds        黄色
            7: 'w',  # unknown       白色
        }.get(var, 'error')

    new_color = []
    new_pts = filter_batch(data_train, label_train)
    for c in new_pts:
        new_color.append(getcolor(c[4]))

    new_pts = new_pts.transpose()
    fig = plt.figure(dpi=120)
    ax = fig.add_subplot(111, projection='3d')
    plt.title('point cloud:{} num:{}'.format(i, len(new_pts[1])))
    # ax.scatter(XYZ[0], XYZ[1], XYZ[2], c=color, marker='.', s=2, linewidth=0, alpha=1, cmap='spectral')
    ax.scatter(new_pts[0], new_pts[1], new_pts[2], c=new_color, marker='.', s=2, linewidth=0, alpha=1, cmap='spectral')

    # ax.set_facecolor((0,0,0))
    ax.axis('scaled')
    # ax.xaxis.set_visible(False)
    # ax.yaxis.set_visible(False)
    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')
    plt.show()
